import torch
from d2l import torch as d2l
import torchvision
from torch import nn
from torch.nn import functional as F


# img = d2l.plt.imread("../data/catdog.jpg")
# h , w = img.shape[:2]
#
# def display_anchors(fmap_w,fmap_h,s):
#     d2l.set_figsize()
#     fmap = torch.zeros((1,3,fmap_h,fmap_w))
#     anchors = d2l.multibox_prior(
#         fmap,sizes=s,ratios=[1,2,0.5])  # 函数返回了有几张图片，锚框的数量，锚框的位置
#     print(anchors[0])
#     bbox_scale = torch.tensor((w,h,w,h))
#     d2l.show_bboxes(d2l.plt.imshow(img).axes,anchors[0]*bbox_scale) # 可以理解最后的anchors[0]是按左上，右下来标记的
#
# display_anchors(1,1,[1.0])
# d2l.plt.show()

def cls_predictor(num_inputs,num_anchors,num_classes):
    return nn.Conv2d(num_inputs,num_anchors*(num_classes+1),kernel_size=3,padding=1)

def bbox_predictor(num_inputs,num_anchors):
    return nn.Conv2d(num_inputs,num_anchors*4,kernel_size=3,padding=1)

def forward(x,block):
    return block(x)

def flatten_pred(pred):
    return torch.flatten(pred.permute(0,2,3,1),start_dim=1)

def concat_preds(preds):
    return torch.cat([flatten_pred(p) for p in preds],dim=1) # 按列拼接

def down_sample_blk(in_channels,out_channels):
    blk = []
    for _ in range(2):
        blk.append(nn.Conv2d(in_channels,out_channels,kernel_size=3,padding=1))
        blk.append(nn.BatchNorm2d(out_channels))
        blk.append(nn.ReLU())
        in_channels = out_channels
    blk.append(nn.MaxPool2d(2))
    return nn.Sequential(*blk)

def base_net():
    blk = []
    num_filters = [3,16,32,64]
    for i in range (len(num_filters)-1):
        blk.append(down_sample_blk(num_filters[i],num_filters[i+1]))
    return nn.Sequential(*blk)

def get_blk(i):
    if i==0:
        blk = base_net()
    elif i==1:
        blk = down_sample_blk(64,128)
    elif i==4:
        blk = nn.AdaptiveAvgPool2d((1,1))
    else:
        blk = down_sample_blk(128,128)
    return blk

def blk_forward(X,blk,size,ratio,cls_predictor,bbox_predictor):
    Y = blk(X)
    anchors = d2l.multibox_prior(Y,sizes=size,ratios=ratio)
    cls_preds = cls_predictor(Y)
    bbox_preds = bbox_predictor(Y)
    return (Y,anchors,cls_preds,bbox_preds)

sizes = [[0.2, 0.272], [0.37, 0.447], [0.54, 0.619], [0.71, 0.79],
         [0.88, 0.961]]
ratios = [[1,2,0.5]] * 5
num_anchors = len(sizes[0]) + len(ratios[0]) - 1




